import importlib
import dataclasses
import os, sys


module_mapping = {
    "gpt-4": {
        "config_name": "GPT4ConfigClass",
        "module_name": "chatgpt_api",
        "class_name": "ChatGPTAPI",
    },
    "gpt-3.5-turbo-16k": {
        "config_name": "GPT35Turbo16kConfigClass",
        "module_name": "chatgpt_api",
        "class_name": "ChatGPTAPI",
    },
    "gpt4all": {
        "config_name": "GPT4ALLConfigClass",
        "module_name": "gpt4all_api",
        "class_name": "GPT4ALLAPI",
    },
}


@dataclasses.dataclass
class GPT4ConfigClass:
    model: str = "gpt-4-browsing"
    api_base: str = "https://api.openai-proxy.com/v1"
    # set up the openai key
    openai_key = os.getenv("OPENAI_KEY", None)
    if openai_key is None:
        print("Your OPENAI_KEY is not set. Please set it in the environment variable.")
    error_wait_time: float = 20
    is_debugging: bool = False


@dataclasses.dataclass
class GPT35Turbo16kConfigClass:
    model: str = "gpt-3.5-turbo-16k"
    api_base: str = "https://api.openai-proxy.com/v1"
    # set up the openai key
    openai_key = os.getenv("OPENAI_KEY", None)
    if openai_key is None:
        print("Your OPENAI_KEY is not set. Please set it in the environment variable.")
    error_wait_time: float = 20
    is_debugging: bool = False


@dataclasses.dataclass
class GPT4ALLConfigClass:
    model: str = "orca-mini-3b.ggmlv3.q4_0.bin"


def dynamic_import(module_name, log_dir) -> object:
    if module_name in module_mapping:
        module_config_name = module_mapping[module_name]["config_name"]
        module_import_name = module_mapping[module_name]["module_name"]
        class_name = module_mapping[module_name]["class_name"]
        module_config = getattr(sys.modules[__name__], module_config_name)
        module_config.log_dir = log_dir

        # import the module
        LLM_module = importlib.import_module(
            "pentestgpt.utils.APIs." + module_import_name
        )
        LLM_class = getattr(LLM_module, class_name)
        # initialize the class
        LLM_class_initialized = LLM_class(module_config)

        return LLM_class_initialized

    else:
        print(
            "Module not found: "
            + module_name
            + ". Falling back to use the default gpt-3.5-turbo-16k"
        )
        # fall back to gpt-3.5-turbo-16k
        LLM_class_initialized = dynamic_import("gpt-3.5-turbo-16k", log_dir)
        return LLM_class_initialized


if __name__ == "__main__":
    # a quick local test
    # load gpt4
    gpt4 = dynamic_import("gpt4all", "logs")
    gpt4.send_new_message("hi")
